import datetime
import requests
from datetime import timezone
import pandas as pd
from utils import df_into_db, read_sql
from quant_researcher.quant.datasource_fetch.crypto_api.glassnode import get_prices
import warnings
warnings.filterwarnings('ignore')


def timestamp_to_utc_str(timestamp, fmt="%Y-%m-%d %H:%M:%S"):
    """
    将时间戳转换为UTC时间字符串
    """
    # 创建UTC时间对象
    utc_dt = datetime.datetime.fromtimestamp(timestamp, tz=timezone.utc)
    # 格式化为字符串
    return utc_dt.strftime(fmt)


def utc_str_to_timestamp(str_time):
    # like"2013-01-01 00:00:00"
    dt = datetime.datetime.strptime(str_time, "%Y-%m-%d %H:%M:%S")
    utc_dt = dt.replace(tzinfo=timezone.utc)
    return int(utc_dt.timestamp())


def update_binance_kline(end_time, frequency):
    assert frequency in ("1h", "1d")
    last_df = read_sql(f"select * from k_line where datasource = 'binance' and frequency = '{frequency}' and "
                       f"type='spot' order by timestamp desc limit 1",
                  db_name="all_history_ohlcvm_coinmarketcap")
    max_db_time = last_df.iloc[0]["timestamp"]
    if frequency == "1d":
        limit = 100
    else:
        limit = 1000
    url = f"https://api.binance.com/api/v3/klines?symbol=BTCUSDT&interval={frequency}&endTime={end_time}&limit={limit}"
    response = requests.get(url)
    data = response.json()
    data = [x[:5] for x in data]
    data = [x for x in data if x[0] >= max_db_time]
    # 检查最后一条数据
    assert last_df.iloc[0]["timestamp"] == data[0][0]
    assert last_df.iloc[0]["open"] == float(data[0][1])
    assert last_df.iloc[0]["high"] == float(data[0][2])
    assert last_df.iloc[0]["low"] == float(data[0][3])
    assert last_df.iloc[0]["close"] == float(data[0][4])
    all_df = pd.DataFrame(data[1:], columns=["datetime", "open", "high", "low", "close"])
    if all_df.empty:
        print(f"binance_{frequency}没有数据需要更新")
        return
    all_df.rename(columns={"datetime": "timestamp"}, inplace=True)
    all_df["datetime"] = all_df["timestamp"].apply(lambda x: timestamp_to_utc_str(x/1000))
    all_df.drop_duplicates(subset=["timestamp"], keep="last", inplace=True)
    all_df["datasource"] = "binance"
    all_df["type"] = "spot"
    all_df["frequency"] = frequency
    df_into_db(all_df, db_name="all_history_ohlcvm_coinmarketcap", table_name="k_line")


def update_binance_kline_v2(end_time, frequency):
    last_df = read_sql(f"select * from k_line where datasource = 'binance' and frequency = '{frequency}' and "
                       f"type='spot' order by timestamp desc limit 1",
                  db_name="all_history_ohlcvm_coinmarketcap")
    max_db_time = last_df.iloc[0]["timestamp"]
    all_datas = []
    if frequency == "1h":
        tmp_end_time = end_time + 3600 * 1000
    elif frequency == "1d":
        tmp_end_time = end_time + (3600*24) * 1000
    elif frequency == "5m":
        tmp_end_time = end_time + 300 * 1000
    else:
        raise Exception("invalid frequency")

    while True:
        limit = 1000  # binance_api最多获取1000条
        url = f"https://api.binance.com/api/v3/klines?symbol=BTCUSDT&interval={frequency}&endTime={tmp_end_time}&limit={limit}"
        response = requests.get(url)
        data = response.json()
        data = [x[:6] + x[7:8] for x in data]
        all_datas.extend(data)
        if int(data[0][0]) <= max_db_time or int(data[-1][0]) <= max_db_time:
            break
        tmp_end_time = min(int(data[0][0]), int(data[-1][0]))
    all_df = pd.DataFrame(all_datas, columns=["timestamp", "open", "high", "low", "close", "vol", "amount"])
    all_df['timestamp'] = all_df['timestamp'].astype('int64')
    all_df['open'] = all_df['open'].astype(float)
    all_df['high'] = all_df['high'].astype(float)
    all_df['low'] = all_df['low'].astype(float)
    all_df['close'] = all_df['close'].astype(float)
    all_df['vol'] = all_df['vol'].astype(float)
    all_df['amount'] = all_df['amount'].astype(float)
    all_df["datetime"] = all_df["timestamp"].apply(lambda x: timestamp_to_utc_str(x/1000))
    all_df.drop_duplicates(subset=["timestamp"], keep="last", inplace=True)
    all_df["datasource"] = "binance"
    all_df["type"] = "spot"
    all_df["frequency"] = frequency
    all_df = all_df[all_df["timestamp"] >= max_db_time]
    all_df = all_df[all_df["timestamp"] <= end_time]
    all_df.sort_values(by="timestamp", ascending=True, inplace=True)
    assert last_df.iloc[0]["timestamp"] == all_df.iloc[0]["timestamp"]
    assert last_df.iloc[0]["open"] == all_df.iloc[0]["open"]
    assert last_df.iloc[0]["high"] == all_df.iloc[0]["high"]
    assert last_df.iloc[0]["low"] == all_df.iloc[0]["low"]
    assert last_df.iloc[0]["close"] == all_df.iloc[0]["close"]
    assert last_df.iloc[0]["vol"] == all_df.iloc[0]["vol"]
    assert last_df.iloc[0]["amount"] == all_df.iloc[0]["amount"]
    all_df = all_df[1:]
    if all_df.empty:
        print(f"binance {frequency} 没有数据需要更新")
        return
    df_into_db(all_df, db_name="all_history_ohlcvm_coinmarketcap", table_name="k_line")


def update_okx_kline(end_time, frequency):
    assert frequency in ("1h", "1d")
    if frequency == "1h":
        bar = "1H"
    else:
        bar = "1Dutc"
    if frequency == "1d":
        limit = 100
    else:
        limit = 1000
    last_df = read_sql(f"select * from k_line where datasource = 'okx' and frequency = '{frequency}' and "
                       f"type='spot' order by timestamp desc limit 1",
                  db_name="all_history_ohlcvm_coinmarketcap")
    max_db_time = last_df.iloc[0]["timestamp"]
    url = f"https://www.okx.com/api/v5/market/history-candles?instId=BTC-USDT&after={end_time}&bar={bar}&limit={limit}"
    response = requests.get(url)
    data = response.json()["data"]
    assert data[0][-1] == "1", "没有未更新好的k线"
    data = [x[:5] for x in data]
    data = [x for x in data if int(x[0]) >= max_db_time]
    # 检查最后一条数据
    assert last_df.iloc[0]["timestamp"] == int(data[-1][0])
    assert last_df.iloc[0]["open"] == float(data[-1][1])
    assert last_df.iloc[0]["high"] == float(data[-1][2])
    assert last_df.iloc[0]["low"] == float(data[-1][3])
    assert last_df.iloc[0]["close"] == float(data[-1][4])
    data = data[:-1]
    if data:
        all_df = pd.DataFrame(data)
        all_df.columns = ["timestamp", "open", "high", "low", "close"]
        all_df['frequency'] = frequency
        all_df['type'] = 'spot'
        all_df["datetime"] = all_df["timestamp"].apply(lambda x: timestamp_to_utc_str(int(x) / 1000))
        all_df["datasource"] = "okx"
        df_into_db(all_df, db_name="all_history_ohlcvm_coinmarketcap", table_name="k_line")
    else:
        print(f"okx {frequency} 没有数据需要更新")


def update_okx_kline_v2(end_time, frequency):
    last_df = read_sql(f"select * from k_line where datasource = 'okx' and frequency = '{frequency}' and "
                       f"type='spot' order by timestamp desc limit 1",
                  db_name="all_history_ohlcvm_coinmarketcap")
    max_db_time = last_df.iloc[0]["timestamp"]
    all_datas = []
    if frequency == "1h":
        bar = "1H"
        tmp_end_time = end_time + 3600 * 1000  # okx不会算end_time的k线
    elif frequency == "1d":
        bar = "1Dutc"
        tmp_end_time = end_time + (3600*24) * 1000  # okx不会算end_time的k线
    elif frequency == "5m":
        bar = "5m"
        tmp_end_time = end_time + 300 * 1000  # okx不会算end_time的k线
    else:
        raise Exception("invalid frequency")

    while True:
        limit = 300  # Okx最大支持300条
        url = f"https://www.okx.com/api/v5/market/history-candles?instId=BTC-USDT&after={tmp_end_time}&bar={bar}&limit={limit}"
        response = requests.get(url)
        data = response.json()["data"]
        assert data[0][-1] == "1", "没有未更新好的k线"
        all_datas.extend(data)
        # 从后往前拿 如果出现小于max_db_time退出
        if int(data[0][0]) <= max_db_time or int(data[-1][0]) <= max_db_time:
            break
        tmp_end_time = min(int(data[0][0]), int(data[-1][0]))
    all_df = pd.DataFrame(all_datas)
    all_df.columns = ["timestamp", "open", "high", "low", "close", "vol", "amount", "amount2", "confirm"]
    assert all_df["amount"].equals(all_df["amount2"]), "btc/usdt报价货币和计价货币都是USDT"
    all_df = all_df[all_df["confirm"] == '1']
    del all_df["confirm"], all_df["amount2"]
    all_df['timestamp'] = all_df['timestamp'].astype('int64')
    all_df['open'] = all_df['open'].astype(float)
    all_df['high'] = all_df['high'].astype(float)
    all_df['low'] = all_df['low'].astype(float)
    all_df['close'] = all_df['close'].astype(float)
    all_df['vol'] = all_df['vol'].astype(float)
    all_df['amount'] = all_df['amount'].astype(float)
    all_df['frequency'] = frequency
    all_df['type'] = 'spot'
    all_df["datetime"] = all_df["timestamp"].apply(lambda x: timestamp_to_utc_str(int(x) / 1000))
    all_df["datasource"] = "okx"
    all_df.sort_values(by="timestamp", ascending=True, inplace=True)
    all_df = all_df[all_df["timestamp"] >= max_db_time]
    all_df = all_df[all_df["timestamp"] <= end_time]
    assert last_df.iloc[0]["timestamp"] == all_df.iloc[0]["timestamp"]
    assert last_df.iloc[0]["open"] == all_df.iloc[0]["open"]
    assert last_df.iloc[0]["high"] == all_df.iloc[0]["high"]
    assert last_df.iloc[0]["low"] == all_df.iloc[0]["low"]
    assert last_df.iloc[0]["close"] == all_df.iloc[0]["close"]
    assert last_df.iloc[0]["vol"] == all_df.iloc[0]["vol"]
    assert last_df.iloc[0]["amount"] == all_df.iloc[0]["amount"]
    all_df = all_df[1:]
    if all_df.empty:
        print(f"okx {frequency} 没有数据需要更新")
        return
    df_into_db(all_df, db_name="all_history_ohlcvm_coinmarketcap", table_name="k_line")


def update_bitstamp_kline(end_time, frequency):
    last_df = read_sql(f"select * from k_line where datasource = 'bitstamp' and frequency = '{frequency}' and "
                       f"type='spot' order by timestamp desc limit 1",
                  db_name="all_history_ohlcvm_coinmarketcap")
    max_db_time = last_df.iloc[0]["timestamp"]
    base_url = "https://www.bitstamp.net/api/v2/ohlc/btcusd/"
    all_data = []
    frequency_map = {"1h": 3600,
                     "5m": 300,
                     "1d": 3600*24}
    start_time = int(max_db_time/1000) - frequency_map[frequency]
    while True:
        params = {
            'step': frequency_map[frequency],
            'limit': 1000,
            'start': start_time,
        }
        response = requests.get(base_url, params=params, timeout=30)
        data = response.json()
        ohlc = data['data']['ohlc']
        all_data.extend(ohlc)
        if len(ohlc) < 1000:
            break
        start_time = int(all_data[-1]["timestamp"])
    df = pd.DataFrame(all_data)
    df["datetime"] = df["timestamp"].apply(lambda x: timestamp_to_utc_str(int(x)))
    df['timestamp'] = df["timestamp"].apply(lambda x: int(x)*1000)
    df['open'] = df['open'].astype(float)
    df['high'] = df['high'].astype(float)
    df['low'] = df['low'].astype(float)
    df['close'] = df['close'].astype(float)
    df.rename(columns={"volume":"vol"}, inplace=True)
    df['vol'] = df['vol'].astype(float)
    df["symbol"] = "BTC"
    df["type"] = "spot"
    df["datasource"] = "bitstamp"
    df["frequency"] = frequency
    print("drop前长度:", len(df))
    df.drop_duplicates(subset=["timestamp"], keep="last", inplace=True)
    print("drop后长度:", len(df))
    df = df[df['timestamp'] <= end_time]
    df = df[df['timestamp'] >= max_db_time]
    assert last_df.iloc[0]["timestamp"] == df.iloc[0]["timestamp"]
    assert last_df.iloc[0]["open"] == df.iloc[0]["open"]
    assert last_df.iloc[0]["high"] == df.iloc[0]["high"]
    assert last_df.iloc[0]["low"] == df.iloc[0]["low"]
    assert last_df.iloc[0]["close"] == df.iloc[0]["close"]
    assert last_df.iloc[0]["vol"] == df.iloc[0]["vol"]
    df = df[1:]
    if df.empty:
        print(f"bitstamp {frequency} 没有数据需要更新")
        return
    df_into_db(df, db_name="all_history_ohlcvm_coinmarketcap", table_name="k_line")


def update_glassnode_kline(frequency, end_date):
    assert frequency in ("1h", "1d")
    last_df = read_sql(f"select * from k_line where datasource = 'glassnode' and frequency = '{frequency}' and "
                       f"type='index' order by timestamp desc limit 1",
                  db_name="all_history_ohlcvm_coinmarketcap")

    if frequency == "1d":
        start_date = last_df.iloc[-1]["datetime"] - datetime.timedelta(days=1)
        start_date = start_date.strftime("%Y-%m-%d")
        interval = "24h"
    else:
        start_date = last_df.iloc[-1]["datetime"] - datetime.timedelta(days=1)
        start_date = start_date.strftime("%Y-%m-%d")
        interval = "1h"
    kline = get_prices(ohlc=True, asset='BTC', start_date=start_date, end_date=end_date, interval=interval)
    kline = kline.reset_index(drop=False)
    kline.rename(columns={"end_date": "datetime"}, inplace=True)
    kline["datetime"] = pd.to_datetime(kline["datetime"])
    end_datetime = datetime.datetime.strptime(end_date, "%Y-%m-%d")
    kline = kline[(kline["datetime"] >= last_df.iloc[-1]["datetime"]) & (kline["datetime"] <= end_datetime)]
    assert round(last_df.iloc[0]["open"], 5) == round(kline.iloc[0]["open"], 5)
    assert round(last_df.iloc[0]["high"], 5) == round(kline.iloc[0]["high"], 5)
    assert round(last_df.iloc[0]["low"], 5) == round(kline.iloc[0]["low"], 5)
    assert round(last_df.iloc[0]["close"], 5) == round(kline.iloc[0]["close"], 5)
    kline = kline[1:]
    if kline.empty:
        print(f"glassnode_{frequency}没有数据需要更新")
        return
    kline["datasource"] = "glassnode"
    kline["frequency"] = frequency
    kline["type"] = "index"
    kline["timestamp"] = kline["datetime"].apply(lambda x: utc_str_to_timestamp(str(x))*1000)
    df_into_db(kline, db_name="all_history_ohlcvm_coinmarketcap", table_name="k_line")


def check_last100_kline(frequency):
    ser_list = []
    if frequency == "5m":
        last_df = read_sql(f"select * from k_line where frequency = '{frequency}' order by timestamp desc limit 300 ",
                           db_name="all_history_ohlcvm_coinmarketcap")
        columns = ["binance","okx", "bitstamp"]
    else:
        last_df = read_sql(f"select * from k_line where frequency = '{frequency}' order by timestamp desc limit 400 ",
                           db_name="all_history_ohlcvm_coinmarketcap")
        columns = ["binance","okx", "bitstamp", "glassnode"]
    for datasource in columns:
        tmp_ser = last_df[last_df["datasource"] == datasource]["datetime"].reset_index(drop=True)
        tmp_ser.name = datasource
        ser_list.append(tmp_ser)
    df = pd.concat(ser_list, axis=1)
    for datasource in columns:
        if datasource != "binance":
            assert df["binance"].equals(df[f"{datasource}"])


def main():
    before_yesterday = datetime.date.today()-datetime.timedelta(days=2)
    end_timestamp = utc_str_to_timestamp(before_yesterday.strftime('%Y-%m-%d %H:%M:%S'))*1000
    print(end_timestamp)
    for frequency in ["1h", "1d", "5m"]:
        if frequency in ("1h", "1d"):
            update_glassnode_kline(frequency, before_yesterday.strftime("%Y-%m-%d"))
        update_bitstamp_kline(end_timestamp, frequency)
        update_binance_kline_v2(end_timestamp, frequency)
        update_okx_kline_v2(end_timestamp, frequency)
        check_last100_kline(frequency)


if __name__ == '__main__':
    main()